import os

import cv2
from tqdm import tqdm


# 将字典写入mat文件
def write_mat(dict, fpath):
    from scipy.io import savemat
    savemat(fpath, dict)


def cap_image(video_root_path, image_save_path):
    """
    读取视频文件并抽帧保存为图片
    @param video_root_path: 视频文件根目录
    @param image_save_path: 图片保存根目录
    """
    print("开始抽帧...")
    for video in tqdm(os.listdir(video_root_path)):
        video_path = os.path.join(video_root_path, video)
        # 船舶类别编号文件夹
        video_name = video.split('.')[0]
        save_path = os.path.join(image_save_path, video_name)
        cap = cv2.VideoCapture(video_path)  # 打开视频文件
        success, frame = cap.read()  # 读取一帧数据
        tracklet = 1
        image_name = 1
        count = 1
        while success:
            # 每8帧取一帧
            if count % 8 == 0:
                # 轨迹文件夹
                tracklet_path = os.path.join(save_path, str(tracklet).zfill(5))
                if not os.path.exists(tracklet_path):
                    os.makedirs(tracklet_path)
                cv2.imwrite(os.path.join(tracklet_path, '{}.jpg'.format(str(image_name).zfill(5))), frame)
                image_name += 1
                # 每8张图片作为一条轨迹
                if image_name % 9 == 0:
                    tracklet += 1
                    image_name = 1
            count += 1
            success, frame = cap.read()  # 继续读取下一帧


def build_dataset(paths, ids, mode="train"):
    """
    用于构建训练集，测试集和基准库，并生成对应的txt文件和mat文件
    Args:
        paths:
        ids:
        mode:

    Returns:

    """
    assert mode in ["train", "test", "baselib"]
    print("\nstart to build {} dataset\n".format(mode))
    txt_path = os.path.join(paths["info_root"], mode + "_name.txt")
    # 存储每个视频的开始帧和结束帧索引，以及对应的vessel_id
    mat_contents = []
    image_nums = 1
    # 测试集的query索引，每艘船舶的第一条轨迹作为query
    query_IDX = []
    num = 1
    for vessel_id in tqdm(ids):
        vessel_path = os.path.join(paths["image_source"], vessel_id)
        first = True
        for track_id in os.listdir(vessel_path):
            track_path = os.path.join(vessel_path, track_id)
            start_num = image_nums
            for image_name in os.listdir(track_path):
                new_image_name = "V{}_T{}_F{}.jpg".format(vessel_id, track_id, image_name.split(".")[0])
                image_path = os.path.join(track_path, image_name)
                new_image_path = os.path.join(paths["image_root"], new_image_name)
                if mode != "baselib":
                    os.system("cp {} {}".format(image_path, new_image_path))
                with open(txt_path, "a") as f:
                    f.write("{}\n".format(new_image_name))
                image_nums += 1
            end_num = image_nums - 1
            mat_contents.append([start_num, end_num, int(vessel_id)])
            if first and mode == "test":
                first = False
                query_IDX.append(num)
            num += 1
    mat_file = {"tracks_{}_info".format(mode): mat_contents}
    mat_file_path = os.path.join(paths["info_root"], "tracks_" + mode + "_info.mat")
    write_mat(mat_file, mat_file_path)
    if mode == "test":
        query_IDX_mat = {"query_IDX": query_IDX}
        query_IDX_mat_path = os.path.join(paths["info_root"], "query_IDX.mat")
        write_mat(query_IDX_mat, query_IDX_mat_path)


def split_dataset(paths):
    """
    切分数据集，分为训练集、测试集和基准库
    Args:
        paths: 字典，包含各种数据集路径

    Returns:

    """
    vessel_num = len(os.listdir(paths["image_source"]))
    # 70%作为训练集，30%作为测试集
    train_num = int(vessel_num * 0.7)
    train_vessel_ids = os.listdir(paths["image_source"])[:train_num]
    test_vessel_ids = os.listdir(paths["image_source"])[train_num:]
    build_dataset(paths, train_vessel_ids, mode="train")
    build_dataset(paths, test_vessel_ids, mode="test")
    build_dataset(paths, os.listdir(paths["image_source"]), mode="baselib")


def main():
    video_root_path = r'Z:\Code_Pile\Programs\DataProcess\Data\2025_data_cleaned\origin'  # 清洗后视频文件根目录
    save_root_path = r'Z:\Code_Pile\Programs\DataProcess\Data\2025_data_cleaned\frame'  # 抽帧图片保存根目录
    save_dataset_root_path = r'Z:\Code_Pile\Programs\DataProcess\Data\2025_data_cleaned\video_reid_dataset'  # 测试集保存根目录
    os.makedirs(save_root_path, exist_ok=True)
    cap_image(video_root_path, save_root_path)  # 抽帧
    paths = {
        "image_source": save_root_path,
        "image_root": os.path.join(save_dataset_root_path, "images"),
        "info_root": os.path.join(save_dataset_root_path, "info"),
    }
    os.makedirs(paths["image_root"], exist_ok=True)
    os.makedirs(paths["info_root"], exist_ok=True)
    # split_dataset(paths)
    build_dataset(paths, os.listdir(paths["image_source"]), mode="test")  # 构建测试集


if __name__ == "__main__":
    main()
